package com.lm.langchain4j.adapter;

import java.util.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 简单内存向量库，存储向量与原始文本片段。
 * 生产环境换成 PGVector、Milvus、RedisVector 等。
 */
public class InMemoryVectorStore {

    private final AtomicLong idSeq = new AtomicLong(1);
    private final Map<Long, double[]> vectors = new HashMap<>();
    private final Map<Long, String> texts = new HashMap<>();

    public long add(double[] vector, String text) {
        long id = idSeq.getAndIncrement();
        vectors.put(id, vector);
        texts.put(id, text);
        return id;
    }

    public List<SearchResult> search(double[] queryVector, int topK) {
        PriorityQueue<SearchResult> pq = new PriorityQueue<>(Comparator.comparingDouble(r -> r.score));
        for (Map.Entry<Long, double[]> e : vectors.entrySet()) {
            double s = cosineSimilarity(queryVector, e.getValue());
            pq.offer(new SearchResult(e.getKey(), texts.get(e.getKey()), s));
            if (pq.size() > topK) pq.poll();
        }
        List<SearchResult> out = new ArrayList<>();
        while (!pq.isEmpty()) out.add(pq.poll());
        Collections.reverse(out);
        return out;
    }

    public static double cosineSimilarity(double[] a, double[] b) {
        double dot = 0, na = 0, nb = 0;
        int n = Math.min(a.length, b.length);
        for (int i = 0; i < n; i++) {
            dot += a[i] * b[i];
            na += a[i] * a[i];
            nb += b[i] * b[i];
        }
        if (na == 0 || nb == 0) return 0;
        return dot / (Math.sqrt(na) * Math.sqrt(nb));
    }

    public record SearchResult(long id, String text, double score) {}
}

